from langchain_mcp_adapters.client import MultiServerMCPClient
import asyncio
from pydantic import SecretStr
from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent, AgentType
from langchain_core.prompts import PromptTemplate

# 创建模型实例（远程调用）
llm = ChatOpenAI(
    model_name="qwen3-max",
    temperature=0,
    openai_api_key=SecretStr("sk-f4fb03bbc29b4f0995b60dec52645af0"),
    openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
    streaming=True,
)

async def create_amap_mcp_client():
    mcp_config = {
        "amap": {
            "url": "https://mcp.amap.com/sse?key=ae16661d9b2ffda189d36044daf97167",
            "transport": "sse",
        }
    }

    client = MultiServerMCPClient(mcp_config)

    tools = await client.get_tools()

    return client, tools

async def create_and_run_agent():
    client, tools = await create_amap_mcp_client()

    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,  # 是否打印中间思考过程
    )

    prompt_template = PromptTemplate.from_template(
        "你是一个智能助手，可以调用高德 MCP 工具。\n\n问题：{input}"
    )

    prompt = prompt_template.format(input="规划北京南站到北京望京SOHO的公交或地铁路线。")

    resp = await agent.ainvoke(prompt)
    print(resp)

    return resp

asyncio.run(create_and_run_agent())